/*
 * @Description: ICP2 匹配模块
 * @Author: imagoodsong
 * @Date: 2020-10-13
 */
#include "lidar_localization/models/registration/icp2_registration.hpp"
#include "pcl/common/transforms.h"
#include "glog/logging.h"

namespace lidar_localization {

ICP2Registration::ICP2Registration(const YAML::Node& node)
    :kd_tree_(new pcl::KdTreeFLANN <CloudData::POINT>()) {
    
    float max_dis = node["max_dis"].as<float>();
    float euclidean_change = node["euclidean_change"].as<float>();
    float euclidean_eps = node["euclidean_eps"].as<float>();
    int max_iter = node["max_iter"].as<int>();

    SetRegistrationParam(max_dis, euclidean_change, euclidean_eps, max_iter);
}

ICP2Registration::ICP2Registration(float max_dis, float euclidean_change, float euclidean_eps, int max_iter)
    :kd_tree_(new pcl::KdTreeFLANN <CloudData::POINT>()) {

    SetRegistrationParam(max_dis, euclidean_change, euclidean_eps, max_iter);
}

bool ICP2Registration::SetRegistrationParam(float max_dis, float euclidean_change, float euclidean_eps, int max_iter) {
    max_dis_ = max_dis;
    euclidean_change_ = euclidean_change;
    euclidean_eps_ = euclidean_eps;
    max_iter_ = max_iter;

    LOG(INFO) << "ICP 的匹配参数为：" << std::endl
              << "max_dis: " << max_dis << ", "
              << "euclidean_change: " << euclidean_change << ", "
              << "euclidean_eps: " << euclidean_eps << ", "
              << "max_iter: " << max_iter 
              << std::endl << std::endl;

    return true;
}

bool ICP2Registration::SetInputTarget(const CloudData::CLOUD_PTR& input_target) {
    target_ = input_target;
    kd_tree_->setInputCloud(input_target);

    return true;
}

bool ICP2Registration::ScanMatch(const CloudData::CLOUD_PTR& input_source, 
                                const Eigen::Matrix4f& predict_pose, 
                                CloudData::CLOUD_PTR& result_cloud_ptr,
                                Eigen::Matrix4f& result_pose) {
	const double min_err2 = euclidean_eps_ * euclidean_eps_;
	const int n_selected_pts = 100;
	const int step = input_source->points.size() / n_selected_pts;
	
    std::vector<Eigen::Vector3f> source_matched;
	source_matched.reserve ( n_selected_pts );
    std::vector<Eigen::Vector3f> target_matched;
	target_matched.reserve ( n_selected_pts );

    Eigen::Matrix3f R = predict_pose.block<3, 3>(0, 0);
    Eigen::Vector3f t = predict_pose.block<3, 1>(0, 3);
	std::vector <int> index (1);
	std::vector <float> squared_distance (1);
	
	double squared_distance_th = max_dis_ * max_dis_;
	double cur_squared_dist = 0.0;
	double last_squared_dist = std::numeric_limits<double>::max();
	
    for ( int n = 0; n < max_iter_; n ++ ) {

		source_matched.clear();
		target_matched.clear();
		
		double sum_squared_dist = 0.0;
		
		for ( size_t i = 0; i < input_source->points.size(); i += step ) {
			
            CloudData::POINT pt_d = input_source->points[i];
            Eigen::Vector3f pt(pt_d.x, pt_d.y, pt_d.z);
            Eigen::Vector3f pt_trans = R * pt  + t;

			pt_d = CloudData::POINT(pt_trans[0], pt_trans[1], pt_trans[2]);
			if (!kd_tree_->nearestKSearch (pt_d, 1, index, squared_distance))
			{
				std::cerr << "ERROR: no points found.\n";
				continue;
			}
			
			if(squared_distance[0] < squared_distance_th)
			{
				sum_squared_dist += squared_distance[0];
				source_matched.push_back(pt);
                pt_d = target_->at(index[0]);
				target_matched.push_back(Eigen::Vector3f(pt_d.x, pt_d.y, pt_d.z));
			}
        }

        Eigen::Vector3f mu_x(0.0, 0.0, 0.0);
		Eigen::Vector3f mu_p(0.0, 0.0, 0.0);
        for(size_t i = 0; i < source_matched.size(); i ++)
		{
			mu_x += target_matched.at(i);
			mu_p += source_matched.at(i);
		}
        mu_x = mu_x / float(target_matched.size());
		mu_p = mu_p / float(source_matched.size());
		
		Eigen::Matrix3f W = Eigen::Matrix3f::Zero();
		for(size_t i = 0; i < source_matched.size(); i ++)
		{
			W += (target_matched.at(i)-mu_x) * ( (source_matched.at(i)-mu_p).transpose() );
		}
		
		Eigen::JacobiSVD<Eigen::Matrix3f> svd(W, Eigen::ComputeFullU | Eigen::ComputeFullV);
		
		R = svd.matrixU() * (svd.matrixV().transpose());
		
		t  = mu_x - R * mu_p;
		
		cur_squared_dist = sum_squared_dist / (double)source_matched.size();
		double squared_dist_change = last_squared_dist -cur_squared_dist;
		
		if(squared_dist_change < euclidean_change_ * euclidean_change_ || cur_squared_dist < min_err2)
		{
			break;
		}
		last_squared_dist = cur_squared_dist;
    }
    result_pose = Eigen::Matrix4f::Identity();
    result_pose.block<3, 3>(0, 0) = R;
    result_pose.block<3, 1>(0, 3) = t;
    pcl::transformPointCloud(*input_source, *result_cloud_ptr, result_pose);
    return true;
}
}